{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Binary Semantic Segmentation using TAO UNET\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\"> "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Take a pretrained resnet18 model and train a ResNet-18 UNet model on the ISBI dataset\n",
    "* Run Inference on the trained model and visualize the inferences\n",
    "* Export the trained model to a .onnx file for deployment to DeepStream\n",
    "* Run inference on the exported .onnx model to verify deployment using TensorRT\n",
    "\n",
    "At the end of this notebook, you will have generated a trained and optimized `UNet` model trained\n",
    "on the ISBI dataset. You may then deploy this model via [Triton](https://github.com/NVIDIA-AI-IOT/tao-toolkit-triton-apps)\n",
    "or [DeepStream](https://developer.nvidia.com/deepstream-sdk).\n",
    "\n",
    "### Table of Contents\n",
    "\n",
    "This notebook shows an example use case of UNet Binary Semantic Segmentation using Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables](#head-0)\n",
    "1. [Installing the TAO launcher](#head-1) <br>\n",
    "2. [Prepare dataset and pre-trained model](#head-2) <br>\n",
    "    1. [Verify downloaded dataset](#head-2-1)\n",
    "    1. [Prepare ISBI data from TIF](#head-2-2)\n",
    "    2.[Visualize the Groundtruth Masks](#head-2-3)\n",
    "    3. [Download pre-trained model](#head-2-4)\n",
    "3. [Provide training specification](#head-3)\n",
    "4. [Run TAO training](#head-4)\n",
    "5. [Evaluate trained models](#head-5)\n",
    "6. [Visualizing Inferences](#head-6)\n",
    "7. [Model Export](#head-7)\n",
    "8. [Verify the deployed model](#head-8)\n",
    "9. [QAT workflow](#head-9)\n",
    "   1. [Convert pruned model to QAT and retrain](#head-9-1)\n",
    "   2. [Evaluate QAT converted model](#head-9-2)\n",
    "   3. [Export QAT trained model to int8](#head-9-3)\n",
    "   4. [Evaluate a QAT trained model using the exported TensorRT engine](#head-9-4)\n",
    "   5. [Inference using QAT engine](#head-9-5)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sample prediction of UNet\n",
    "<img align=\"center\" src=\"https://github.com/vpraveen-nv/model_card_images/blob/main/cv/notebook/unet/sample.jpg?raw=true\" width=\"300\" height=\"300\">"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "The following notebook requires the user to set an env variable called the `$LOCAL_PROJECT_DIR` as the path to the users workspace. Please note that the dataset to run this notebook is expected to reside in the `$LOCAL_PROJECT_DIR/data/isbi`, while the TAO experiment generated collaterals will be output to `$LOCAL_PROJECT_DIR/unet`. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells.\n",
    "\n",
    "*Note: Please make sure to remove any stray artifacts/files from the `$USER_EXPERIMENT_DIR` or `$DATA_DOWNLOAD_DIR` paths as mentioned below, that may have been generated from previous experiments. Having checkpoint files etc may interfere with creating a training graph for a new experiment.*\n",
    "\n",
    "*Note: This notebook currently is by default set up to run training using 1 GPU. To use more GPU's please update the env variable `$NUM_GPUS` accordingly*\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up env variables for cleaner command line commands.\n",
    "import os\n",
    "\n",
    "%set_env GPU_INDEX=\"0\"\n",
    "%env NUM_GPUS=1\n",
    "%set_env USER_EXPERIMENT_DIR=/workspace/tao-experiments/unet\n",
    "%set_env DATA_DOWNLOAD_DIR=/workspace/tao-experiments/data\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=~/tao-samples/unet\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "# The dataset expected to be present in $LOCAL_PROJECT_DIR/data, while the results for the steps\n",
    "# in this notebook will be stored at $LOCAL_PROJECT_DIR/unet\n",
    "# !PLEASE MAKE SURE TO UPDATE THIS PATH!.\n",
    "os.environ[\"LOCAL_PROJECT_DIR\"] = FIXME\n",
    "\n",
    "# !PLEASE MAKE SURE TO UPDATE THIS PATH!.\n",
    "# Point to the 'deps' folder in samples from where you are launching notebook inside unet folder.\n",
    "%env PROJECT_DIR=/workspace/iva/ngc-collaterals/cv/samples\n",
    "\n",
    "os.environ[\"LOCAL_DATA_DIR\"] = os.path.join(\n",
    "    os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()),\n",
    "    \"data\"\n",
    ")\n",
    "os.environ[\"LOCAL_EXPERIMENT_DIR\"] = os.path.join(\n",
    "    os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()),\n",
    "    \"unet\"\n",
    ")\n",
    "\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"LOCAL_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")\n",
    "%set_env SPECS_DIR=/workspace/tao-experiments/unet/specs\n",
    "\n",
    "! ls -l $LOCAL_DATA_DIR\n",
    "!ls -rlt $LOCAL_SPECS_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cell below maps the project directory on your local host to a workspace directory in the TAO docker instance, so that the data and the results are mapped from in and out of the docker. For more information please refer to the [launcher instance](https://docs.nvidia.com/tao/tao-toolkit/text/tao_launcher.html) in the user guide.\n",
    "\n",
    "When running this cell on AWS, update the drive_map entry with the dictionary defined below, so that you don't have permission issues when writing data into folders created by the TAO docker.\n",
    "\n",
    "```json\n",
    "drive_map = {\n",
    "    \"Mounts\": [\n",
    "            # Mapping the data directory\n",
    "            {\n",
    "                \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "                \"destination\": \"/workspace/tao-experiments\"\n",
    "            },\n",
    "            # Mapping the specs directory.\n",
    "            {\n",
    "                \"source\": os.environ[\"LOCAL_SPECS_DIR\"],\n",
    "                \"destination\": os.environ[\"SPECS_DIR\"]\n",
    "            },\n",
    "        ],\n",
    "    \"DockerOptions\": {\n",
    "        \"user\": \"{}:{}\".format(os.getuid(), os.getgid())\n",
    "    }\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "\n",
    "# Define the dictionary with the mapped drives\n",
    "drive_map = {\n",
    "    \"Mounts\": [\n",
    "        # Mapping the data directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "        # Mapping the specs directory.\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_SPECS_DIR\"],\n",
    "            \"destination\": os.environ[\"SPECS_DIR\"]\n",
    "        },\n",
    "    ]\n",
    "}\n",
    "\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(drive_map, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install requirement\n",
    "!pip3 install Cython==0.29.36\n",
    "!pip3 install -r $PROJECT_DIR/deps/requirements-pip.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in PyPI. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare dataset and pre-trained model <a class=\"anchor\" id=\"head-1\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will be using the `ISBI Challenge: Segmentation of neuronal structures in EM stacks` dataset for the binary segmentation tutorial. Please access the open source repo here https://github.com/alexklibisz/isbi-2012/tree/master/data to download the data. The data is in .tif format. Copy the train-labels.tif, train-volume.tif, test-volume.tif files to `$LOCAL_DATA_DIR/isbi`. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Verify downloaded dataset <a class=\"anchor\" id=\"head-1-1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the .tiff files are present\n",
    "!mkdir -p $LOCAL_DATA_DIR\n",
    "!ls -l $LOCAL_DATA_DIR\n",
    "!if [ ! -f $LOCAL_DATA_DIR/isbi/test-volume.tif ]; then echo 'test-volume.tif file not found, please download.'; else echo 'Found test-volume.tif file.';fi\n",
    "!if [ ! -f $LOCAL_DATA_DIR/isbi/test-volume.tif ]; then echo 'train-labels file not found, please download.'; else echo 'Found train-labels.tif file.';fi\n",
    "!if [ ! -f $LOCAL_DATA_DIR/isbi/test-volume.tif ]; then echo 'train-volume.tif file not found, please download.'; else echo 'Found train-volume.tif file.';fi"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### B.  Prepare ISBI data from TIF <a class=\"anchor\" id=\"head-1-2\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prepare the images and masks from the .tif files by running the following script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!bash prepare_data.sh $LOCAL_DATA_DIR/isbi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!if [ -f $DATA_DOWNLOAD_DIR/isbi/images/train/image_0.png ]; then echo 'Data prepared successfully !'; else echo 'Please verify Data Preparation.';fi\n",
    "# Checking the files\n",
    "!ls -l $LOCAL_DATA_DIR/isbi\n",
    "!ls -l $LOCAL_DATA_DIR/isbi/images\n",
    "!ls -l $LOCAL_DATA_DIR/isbi/masks\n",
    "!ls -l $LOCAL_DATA_DIR/isbi/images/train"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### C.  Visualize the Groundtruth Masks <a class=\"anchor\" id=\"head-1-3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run this cell to visualize the training data and the masks\n",
    "\n",
    "!python vis_annotation_isbi.py -i $LOCAL_DATA_DIR/isbi/images/train -m $LOCAL_DATA_DIR/isbi/masks/train -o $LOCAL_EXPERIMENT_DIR/isbi_experiment_unpruned/vis_gt --num_classes 2 --num_images 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple grid visualizer\n",
    "%matplotlib inline\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from math import ceil\n",
    "valid_image_ext = ['.jpg', '.png']\n",
    "\n",
    "def visualize_images(image_dir, result_dir, num_cols=4, num_images=10):\n",
    "    output_path = os.path.join(os.environ['LOCAL_EXPERIMENT_DIR'], result_dir, image_dir)\n",
    "    num_rows = int(ceil(float(num_images) / float(num_cols)))\n",
    "    f, axarr = plt.subplots(num_rows, num_cols, figsize=[80,30])\n",
    "    f.tight_layout()\n",
    "    a = [os.path.join(output_path, image) for image in os.listdir(output_path) \n",
    "         if os.path.splitext(image)[1].lower() in valid_image_ext]\n",
    "    for idx, img_path in enumerate(a[:num_images]):\n",
    "        col_id = idx % num_cols\n",
    "        row_id = idx // num_cols\n",
    "        img = plt.imread(img_path)\n",
    "        axarr[row_id, col_id].imshow(img) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images.\n",
    "OUTPUT_PATH = 'vis_gt' # relative path from $USER_EXPERIMENT_DIR.\n",
    "COLS = 2 # number of columns in the visualizer grid.\n",
    "IMAGES = 4 # number of images to visualize.\n",
    "\n",
    "visualize_images(OUTPUT_PATH, \"isbi_experiment_unpruned\", num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### D. Download pre-trained model <a class=\"anchor\" id=\"head-1-4\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use NGC CLI to get the pre-trained models. For more details, go to ngc.nvidia.com and click the SETUP on the navigation bar. \n",
    "\n",
    "*Note: When using vanilla_unet as arch for binary segmentation, pre-trained model section can be skipped. Pre-trained weights are available only for Resnet/ VGG templates*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $LOCAL_PROJECT_DIR/ngccli\n",
    "\n",
    "# Remove any previously existing CLI installations\n",
    "!rm -rf $LOCAL_PROJECT_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $LOCAL_PROJECT_DIR/ngccli\n",
    "!unzip -u \"$LOCAL_PROJECT_DIR/ngccli/$CLI\" -d $LOCAL_PROJECT_DIR/ngccli/\n",
    "!rm $LOCAL_PROJECT_DIR/ngccli/*.zip \n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"LOCAL_PROJECT_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ngc registry model list nvidia/tao/pretrained_semantic_segmentation:*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/pretrained_resnet18/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pull pretrained model from NGC\n",
    "!ngc registry model download-version nvidia/tao/pretrained_semantic_segmentation:resnet18 --dest $LOCAL_EXPERIMENT_DIR/pretrained_resnet18"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Check that model is downloaded into dir.\")\n",
    "!ls -l $LOCAL_EXPERIMENT_DIR/pretrained_resnet18/pretrained_semantic_segmentation_vresnet18"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide ISBI training specification <a class=\"anchor\" id=\"head-3\"></a>\n",
    "\n",
    "* Images and Masks path\n",
    "    * In order to use the newly generated images, masks folder update the dataset_config parameter in the spec file at `$SPECS_DIR/unet_train_resnet_unet_isbi.txt` \n",
    "    * Update the train, val images and masks paths. The test only requires the images path. \n",
    "* Pre-trained models\n",
    "* Augmentation parameters for on the fly data augmentation\n",
    "* Other training (hyper-)parameters such as batch size, number of epochs, learning rate etc.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!cat $LOCAL_SPECS_DIR/unet_train_resnet_unet_isbi.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run TAO training <a class=\"anchor\" id=\"head-4\"></a>\n",
    "* Provide the sample spec file and the output directory location for models\n",
    "* WARNING: training will take several hours or one day to complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(\"For multi-GPU, change --gpus based on your machine.\")\n",
    "!tao model unet train --gpus $NUM_GPUS \\\n",
    "                      --gpu_index $GPU_INDEX \\\n",
    "                      -e $SPECS_DIR/unet_train_resnet_unet_isbi.txt \\\n",
    "                      -r $USER_EXPERIMENT_DIR/isbi_experiment_unpruned \\\n",
    "                      -m $USER_EXPERIMENT_DIR/pretrained_resnet18/pretrained_semantic_segmentation_vresnet18/resnet_18.hdf5 \\\n",
    "                      -n model_isbi"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unet supports restarting from checkpoint. In case, the training job is killed prematurely, you may resume training from the closest checkpoint by simply re-running the same command line. Please do make sure to use the same number of GPUs when restarting the training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Model for every epoch at checkpoint_interval mentioned in the spec file:')\n",
    "print('---------------------')\n",
    "!ls -ltrh $LOCAL_EXPERIMENT_DIR/isbi_experiment_unpruned/\n",
    "!ls -ltrh $LOCAL_EXPERIMENT_DIR/isbi_experiment_unpruned/weights"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate trained models <a class=\"anchor\" id=\"head-5\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last step model saved in the `$USER_EXPERIMENT_DIR/isbi_experiment_unpruned/weights` dir is used for evaluation/ inference/ export. The evaluation also creates `$LOCAL_EXPERIMENT_DIR/isbi_experiment_unpruned/results_tlt.json`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tao model unet evaluate --gpu_index $GPU_INDEX -e $SPECS_DIR/unet_train_resnet_unet_isbi.txt \\\n",
    "                 -m $USER_EXPERIMENT_DIR/isbi_experiment_unpruned/weights/model_isbi.tlt \\\n",
    "                 -o $USER_EXPERIMENT_DIR/isbi_experiment_unpruned/ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $LOCAL_EXPERIMENT_DIR/isbi_experiment_unpruned/results_tlt.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Visualizing Inferences <a class=\"anchor\" id=\"head-6\"></a>\n",
    "In this section, we run the UNet inference tool to generate inferences on the trained models and print the results. \n",
    "\n",
    "The following cell will run inference for segmentation and visualize masks for the images in test. The resulting visualized images will be saved in the `vis_overlay_tlt` folder and label PNG masks in `mask_labels_tlt` in the path provided to `-o` argument.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model unet inference --gpu_index $GPU_INDEX -e $SPECS_DIR/unet_train_resnet_unet_isbi.txt \\\n",
    "                 -m $USER_EXPERIMENT_DIR/isbi_experiment_unpruned/weights/model_isbi.tlt \\\n",
    "                 -o $USER_EXPERIMENT_DIR/isbi_experiment_unpruned/ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images.\n",
    "OUTPUT_PATH = 'vis_overlay_tlt' # relative path from $USER_EXPERIMENT_DIR.\n",
    "COLS = 2 # number of columns in the visualizer grid.\n",
    "IMAGES = 4 # number of images to visualize.\n",
    "\n",
    "visualize_images(OUTPUT_PATH, \"isbi_experiment_unpruned\", num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Prune the trained model <a class=\"anchor\" id=\"head-6\"></a>\n",
    "* Specify pre-trained model\n",
    "* Equalization criterion (`Applicable for resnets`)\n",
    "* Threshold for pruning.\n",
    "* Output directory to store the model\n",
    "\n",
    "*Usually, you just need to adjust `-pth` (threshold) for accuracy and model size trade off. Higher `pth` gives you smaller model (and thus higher inference speed) but worse accuracy. The threshold to use is dependent on the dataset. A pth value `5.2e-6` is just a start point. If the retrain accuracy is good, you can increase this value to get smaller models. Otherwise, lower this value to get better accuracy.*\n",
    "\n",
    "* For some internal studies, we have noticed that a pth value of 0.1 is a good starting point for unet models trained on larger datasets. A larger regularization value in the first round of training will result in smaller models while pruning. Hence regularization while training and pth are hyper-parameters that needs to be tuned.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an output directory if it doesn't exist.\n",
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/isbi_experiment_pruned"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model unet prune \\\n",
    "                  -e $SPECS_DIR/unet_train_resnet_unet_isbi.txt \\\n",
    "                  -m $USER_EXPERIMENT_DIR/isbi_experiment_unpruned/weights/model_isbi.tlt \\\n",
    "                  -o $USER_EXPERIMENT_DIR/isbi_experiment_pruned/model_isbi_pruned.tlt \\\n",
    "                  -eq union \\\n",
    "                  -pth 0.6 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -rlt $LOCAL_EXPERIMENT_DIR/isbi_experiment_pruned/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Retrain the pruned model <a class=\"anchor\" id=\"head-7\"></a>\n",
    "* Model needs to be re-trained to bring back accuracy after pruning\n",
    "* Specify re-training specification with pretrained weights as pruned model.\n",
    "\n",
    "*Note: For retraining, please set the `load_graph` option to `true` in the model_config to load the pruned model graph. Also, if after retraining, the model shows some decrease in MIOU, it could be that the originally trained model was pruned a little too much. Please try reducing the pruning threshold (thereby reducing the pruning ratio) and use the new model to retrain.*\n",
    "\n",
    "*Note: Ensure to provide a different folder for saving results of retraining from the folder where pruned model is saved."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Printing the retrain experiment file. \n",
    "# Note: We have updated the experiment file to include the \n",
    "# newly pruned model as a pretrained weights and, the\n",
    "# load_graph option is set to true \n",
    "!cat $LOCAL_SPECS_DIR/unet_train_resnet_unet_isbi_retrain.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retraining using the pruned model as pretrained weights \n",
    "!tao model unet train --gpus $NUM_GPUS \\\n",
    "                      --gpu_index $GPU_INDEX \\\n",
    "                      -e $SPECS_DIR/unet_train_resnet_unet_isbi_retrain.txt \\\n",
    "                      -r $USER_EXPERIMENT_DIR/isbi_experiment_retrain \\\n",
    "                      -m $USER_EXPERIMENT_DIR/isbi_experiment_pruned/model_isbi_pruned.tlt \\\n",
    "                      -n model_isbi_retrained "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Listing the newly retrained model.\n",
    "!ls -rlt $LOCAL_EXPERIMENT_DIR/isbi_experiment_retrain/weights"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Evaluate the retrained model <a class=\"anchor\" id=\"head-8\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This section evaluates the pruned and retrained model, using the `evaluate` command."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model unet evaluate --gpu_index $GPU_INDEX -e $SPECS_DIR/unet_train_resnet_unet_isbi_retrain.txt \\\n",
    "                 -m $USER_EXPERIMENT_DIR/isbi_experiment_retrain/weights/model_isbi_retrained.tlt \\\n",
    "                 -o $USER_EXPERIMENT_DIR/isbi_experiment_retrain/ "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Model Export <a class=\"anchor\" id=\"head-7\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# tao <task> export will fail if .onnx already exists. So we clear the export folder before tao <task> export\n",
    "!rm -rf $LOCAL_EXPERIMENT_DIR/export\n",
    "# Generate .onnx file using tao container\n",
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/export \n",
    "\n",
    "!tao model unet export --gpu_index $GPU_INDEX -m $USER_EXPERIMENT_DIR/isbi_experiment_retrain/weights/model_isbi_retrained.tlt \\\n",
    "               -e $SPECS_DIR/unet_train_resnet_unet_isbi_retrain.txt \\\n",
    "               --gen_ds_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check if onnx model is correctly saved.\n",
    "!ls -l $LOCAL_EXPERIMENT_DIR/isbi_experiment_retrain/weights/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the `tao deploy` container, you can generate a TensorRT engine and verify the correctness of the generated through evaluate and inference. \n",
    "\n",
    "The `tao deploy` produces optimized tensorrt engines for the platform that it resides on. Therefore, to get maximum performance, please run `tao deploy` command which will instantiate a deploy container, with the exported `.onnx` file on your target device. The `tao deploy` container only works for x86, with discrete NVIDIA GPU's. \n",
    "\n",
    "If you choose to integrate your model into deepstream directly, please refer to [deepstream dev guide](https://docs.nvidia.com/tao/tao-toolkit/text/ds_tao/deepstream_tao_integration.html) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert to TensorRT engine(FP32).\n",
    "!tao deploy unet gen_trt_engine --gpu_index $GPU_INDEX \\\n",
    "                                -m $USER_EXPERIMENT_DIR/isbi_experiment_retrain/weights/model_isbi_retrained.onnx \\\n",
    "                                -e $SPECS_DIR/unet_train_resnet_unet_isbi_retrain.txt \\\n",
    "                                -r $USER_EXPERIMENT_DIR/export \\\n",
    "                                --data_type fp32 \\\n",
    "                                --engine_file $USER_EXPERIMENT_DIR/export/trtfp32.isbi.retrained.engine \\\n",
    "                                --max_batch_size 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verify the tensorrt engine accuracy on the validation dataset through tao deploy. This generates results text file \n",
    "# results.json in the directory mentioned to the -r argument \n",
    "\n",
    "!tao deploy unet evaluate --gpu_index $GPU_INDEX -e $SPECS_DIR/unet_train_resnet_unet_isbi_retrain.txt \\\n",
    "                          -m $USER_EXPERIMENT_DIR/export/trtfp32.isbi.retrained.engine \\\n",
    "                          -r $USER_EXPERIMENT_DIR/isbi_experiment_retrain/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $LOCAL_EXPERIMENT_DIR/isbi_experiment_retrain/results.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert to TensorRT engine(INT8).\n",
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/export \n",
    "!tao deploy unet gen_trt_engine --gpu_index $GPU_INDEX -m $USER_EXPERIMENT_DIR/isbi_experiment_retrain/weights/model_isbi_retrained.onnx \\\n",
    "                                -e $SPECS_DIR/unet_train_resnet_unet_isbi_retrain.txt \\\n",
    "                                -r $USER_EXPERIMENT_DIR/export \\\n",
    "                                --data_type int8 \\\n",
    "                                --engine_file $USER_EXPERIMENT_DIR/export/int8.isbi.retrained.engine \\\n",
    "                                --data_type int8 \\\n",
    "                                --cal_data_file $USER_EXPERIMENT_DIR/export/isbi_cal_data_file.txt \\\n",
    "                                --cal_cache_file $USER_EXPERIMENT_DIR/export/isbi_cal.bin \\\n",
    "                                --cal_image_dir $DATA_DOWNLOAD_DIR/isbi/images/val \\\n",
    "                                --max_batch_size 3 \\\n",
    "                                --batch_size 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported model:')\n",
    "print('------------')\n",
    "!ls -lh $LOCAL_EXPERIMENT_DIR/export"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Verify the deployed model <a class=\"anchor\" id=\"head-8\"></a>\n",
    "\n",
    "Verify the converted engine. The resulting TRT inference images will be saved in `vis_overlay` folder and PNG masks in `mask_labels`  in the path provided to `-o` argument.\n",
    "\n",
    "*Note: The TensorRT evaluation and inference must be run with a batch size smaller than `max_batch_size` that was used while generating engine.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Running inference on test dataset through tao deploy\n",
    "!tao deploy unet inference --gpu_index $GPU_INDEX -e $SPECS_DIR/unet_train_resnet_unet_isbi_retrain.txt \\\n",
    "                           -m $USER_EXPERIMENT_DIR/export/int8.isbi.retrained.engine \\\n",
    "                           -r $USER_EXPERIMENT_DIR/isbi_experiment_retrain/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Running inference on test dataset through tao deploy\n",
    "!tao deploy unet evaluate --gpu_index $GPU_INDEX -e $SPECS_DIR/unet_train_resnet_unet_isbi_retrain.txt \\\n",
    "                          -m $USER_EXPERIMENT_DIR/export/int8.isbi.retrained.engine \\\n",
    "                          -r $USER_EXPERIMENT_DIR/isbi_experiment_retrain/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images.\n",
    "OUTPUT_PATH = 'vis_overlay' # relative path from $USER_EXPERIMENT_DIR.\n",
    "COLS = 2 # number of columns in the visualizer grid.\n",
    "IMAGES = 4 # number of images to visualize.\n",
    "\n",
    "visualize_images(OUTPUT_PATH, \"isbi_experiment_retrain\", num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10. QAT workflow <a class=\"anchor\" id=\"head-12\"></a>\n",
    "This section delves into the newly enabled Quantization Aware Training feature with UNet. The workflow defined below converts a pruned model from section [5](#head-5) to enable QAT and retrain this model to while accounting the noise introduced due to quantization in the forward pass. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Convert pruned model to QAT and retrain <a class=\"anchor\" id=\"head-12-1\"></a>\n",
    "All UNet models, unpruned and pruned models can be converted to QAT models by setting the `enable_qat` parameter in the `training_config` component of the spec file to `true`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Printing the retrain experiment file. \n",
    "# Note: We have updated the experiment file to convert the\n",
    "# pretrained model to qat mode by setting the enable_qat\n",
    "# parameter.\n",
    "!cat $LOCAL_SPECS_DIR/unet_train_resnet_unet_isbi_retrain_qat.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model unet train -e $SPECS_DIR/unet_train_resnet_unet_isbi_retrain_qat.txt \\\n",
    "                      -r $USER_EXPERIMENT_DIR/experiment_dir_retrain_qat \\\n",
    "                      -m $USER_EXPERIMENT_DIR/isbi_experiment_pruned/model_isbi_pruned.tlt \\\n",
    "                      -n resnet18_unet_pruned_qat \\\n",
    "                      --gpus $NUM_GPUS \\\n",
    "                      --gpu_index $GPU_INDEX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -rlt $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain_qat/weights"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### B. Evaluate QAT converted model <a class=\"anchor\" id=\"head-12-2\"></a>\n",
    "This section evaluates a QAT enabled pruned retrained model. The MIOU of this model should be comparable to that of the pruned retrained model without QAT. However, due to quantization, it is possible sometimes to see a drop in the MIOU value for certain datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model unet evaluate -e $SPECS_DIR/unet_train_resnet_unet_isbi_retrain_qat.txt \\\n",
    "                           -m $USER_EXPERIMENT_DIR/experiment_dir_retrain_qat/weights/resnet18_unet_pruned_qat.tlt \\\n",
    "                           -o $USER_EXPERIMENT_DIR/isbi_experiment_pruned_qat/ "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### C. Export QAT trained model to int8 <a class=\"anchor\" id=\"head-12-3\"></a>\n",
    "Export a QAT trained model to TensorRT parsable model. This command generates an .onnx file from the trained model and the serializes corresponding int8 scales as a TRT readable calibration cache file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment the below lines and run for QA to run Int8 export also in the next cell \n",
    "!rm -rf $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain_qat/weights/*.onnx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model unet export --gpu_index $GPU_INDEX -m $USER_EXPERIMENT_DIR/experiment_dir_retrain_qat/weights/resnet18_unet_pruned_qat.tlt \\\n",
    "                       -e $SPECS_DIR/unet_train_resnet_unet_isbi_retrain_qat.txt \\\n",
    "                       --cal_json_file $USER_EXPERIMENT_DIR/export/isbi_cal_qat.json \\\n",
    "                       --gen_ds_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert to TensorRT engine\n",
    "!tao deploy unet gen_trt_engine --gpu_index $GPU_INDEX -m $USER_EXPERIMENT_DIR/experiment_dir_retrain_qat/weights/resnet18_unet_pruned_qat.onnx \\\n",
    "                                -e $SPECS_DIR/unet_train_resnet_unet_isbi_retrain_qat.txt \\\n",
    "                                --engine_file $USER_EXPERIMENT_DIR/export/int8.isbi.retrained_qat.engine \\\n",
    "                                -r $USER_EXPERIMENT_DIR/export \\\n",
    "                                --data_type int8 \\\n",
    "                                --cal_json_file $USER_EXPERIMENT_DIR/export/isbi_cal_qat.json \\\n",
    "                                --cal_cache_file $USER_EXPERIMENT_DIR/export/isbi_cal_qat.bin \\\n",
    "                                --max_batch_size 3 \\\n",
    "                                --batch_size 1"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### D. Evaluate a QAT trained model using the exported TensorRT engine <a class=\"anchor\" id=\"head-12-4\"></a>\n",
    "This section evaluates a QAT enabled pruned retrained model using the TensorRT int8 engine that was exported in [Section C](#head-12-3). Please note that there maybe a slight difference (~0.1-0.5%) in the MIOU from [Section B](#head-12-2), owing to some differences in the implementation of quantization in TensorRT.\n",
    "\n",
    "*Note: The TensorRT evaluator might be slightly slower than the TAO evaluator here, because the evaluation dataloader is pinned to the CPU to avoid any clashes between TensorRT and TAO instances in the GPU. Please note that this tool was not intended and has not been developed for profiling the model. It is just a means to qualitatively analyse the model.*\n",
    "\n",
    "*Please use native TensorRT or DeepStream for the most optimized inferences.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Running inference on test dataset through tao deploy\n",
    "!tao deploy unet evaluate --gpu_index $GPU_INDEX -e $SPECS_DIR/unet_train_resnet_unet_isbi_retrain_qat.txt \\\n",
    "                          -m $USER_EXPERIMENT_DIR/export/int8.isbi.retrained_qat.engine \\\n",
    "                          -r $USER_EXPERIMENT_DIR/isbi_experiment_pruned_qat/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### E. Inference using QAT engine <a class=\"anchor\" id=\"head-12-5\"></a>\n",
    "Run inference and visualize detections on test images, using the exported TensorRT engine from [Section C](#head-12-3)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Running inference on test dataset through tao deploy\n",
    "!tao deploy unet inference --gpu_index $GPU_INDEX -e $SPECS_DIR/unet_train_resnet_unet_isbi_retrain_qat.txt \\\n",
    "                           -m $USER_EXPERIMENT_DIR/export/int8.isbi.retrained_qat.engine \\\n",
    "                           -r $USER_EXPERIMENT_DIR/isbi_experiment_pruned_qat/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images.\n",
    "OUTPUT_PATH = 'vis_overlay' # relative path from $USER_EXPERIMENT_DIR.\n",
    "COLS = 2 # number of columns in the visualizer grid.\n",
    "IMAGES = 4 # number of images to visualize.\n",
    "\n",
    "visualize_images(OUTPUT_PATH, \"isbi_experiment_pruned_qat\", num_cols=COLS, num_images=IMAGES)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
